import torch
from torch.utils.data import Dataset
import numpy as np
import os
import json

class NCPDDataset(Dataset):
    def __init__(self, data_root, config, split='train'):
        self.data_root = data_root
        self.config = config
        self.split = split
        
        # 加载数据列表
        with open(os.path.join(data_root, f'{split}_list.json'), 'r') as f:
            self.data_list = json.load(f)
        
        self.history_frames = config['training']['history_frames']
        self.future_frames = config['training']['future_frames']
    
    def __len__(self):
        return len(self.data_list)
    
    def __getitem__(self, idx):
        data_id = self.data_list[idx]
        
        # 加载红外数据
        infrared_path = os.path.join(self.data_root, 'infrared', f'{data_id}.npy')
        infrared_data = np.load(infrared_path)  # [seq_len, 1, H, W]
        
        # 加载雷达数据
        radar_path = os.path.join(self.data_root, 'radar', f'{data_id}.npy')
        radar_data = np.load(radar_path)  # [seq_len, num_points, 4]
        
        # 加载标签
        label_path = os.path.join(self.data_root, 'labels', f'{data_id}.json')
        with open(label_path, 'r') as f:
            labels = json.load(f)
        
        # 转换为tensor
        infrared_tensor = torch.FloatTensor(infrared_data)
        radar_tensor = torch.FloatTensor(radar_data)
        
        # 处理标签
        intent_labels = torch.FloatTensor(labels['intent'])
        trajectory_labels = torch.FloatTensor(labels['trajectory'])
        
        # 上下文信息
        context = {
            'light_intensity': torch.FloatTensor([labels['light_intensity']]),
            'signal_state': torch.FloatTensor([labels['signal_state']]),
            'in_crosswalk': torch.FloatTensor([labels['in_crosswalk']]),
            'time_of_day': torch.FloatTensor(self.encode_time(labels['timestamp']))
        }
        
        return (infrared_tensor, radar_tensor, context, 
                {'intent': intent_labels, 'trajectory': trajectory_labels})
    
    def encode_time(self, timestamp):
        """将时间编码为周期性特征"""
        # 简化的时间编码
        hour = timestamp % 24
        sin_hour = np.sin(2 * np.pi * hour / 24)
        cos_hour = np.cos(2 * np.pi * hour / 24)
        return [sin_hour, cos_hour]